package org.hepeng.workx.jdbc;

import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author he peng
 */
public class DataSourceRouteContext extends ConcurrentHashMap<String, Object> {

    public static final String DATA_SOURCE_ROUTE_ANNOTATION_KEY = "dataSourceRouteAnnotation";

    private static final ThreadLocal<DataSourceRouteContext> THREAD_LOCAL = new InheritableThreadLocal<DataSourceRouteContext>() {
        @Override
        protected DataSourceRouteContext initialValue() {
            return new DataSourceRouteContext();
        }
    };

    public static void setContext(DataSourceRouteContext context) {
        THREAD_LOCAL.set(context);
    }

    public static DataSourceRouteContext getContext() {
        return THREAD_LOCAL.get();
    }

    public static void close() {
        THREAD_LOCAL.remove();
    }

    public DataSourceRoute getDataSourceRoute() {
        Object obj = get(DATA_SOURCE_ROUTE_ANNOTATION_KEY);
        return Objects.nonNull(obj) ? (DataSourceRoute) obj : null;
    }
}
